#include "types.h"
#include "console.h"
#include "spinlock.h"
#include "sleeplock.h"
#include "fs.h"
#include "buf.h"
#include "string.h"
#include "log.h"
#include "initcode.h"
#include "cop.h"
#include "rand.h"

// #define LOG_TRACE

/* Simple logging that allows concurrent FS system calls.
 *
 * A log transaction contains the updates of multiple FS system
 * calls. The logging system only commits when there are
 * no FS system calls active. Thus there is never
 * any reasoning required about whether a commit might
 * write an uncommitted system call's updates to disk.
 *
 * A system call should call begin_op()/end_op() to mark
 * its start and end. Usually begin_op() just increments
 * the count of in-progress FS system calls and returns.
 * But if it thinks the log is close to running out, it
 * sleeps until the last outstanding end_op() commits.
 *
 * The log is a physical re-do log containing disk blocks.
 * The on-disk log format:
 *   header block, containing block #s for block A, B, C, ...
 *   block A
 *   block B
 *   block C
 *   ...
 * Log appends are synchronous.
 */

/*
 * Contents of the header block, used for both the on-disk header block
 * and to keep track in memory of logged block# before commit.
 */
typedef struct {
    int n;
    int block[LOGSIZE];
} logheader_t;

typedef struct {
    struct spinlock lock;
    int start;
    int size;
    int outstanding;    // How many FS sys calls are executing.
    int committing;     // In commit(), please wait.
    int dev;
    int seq, ack;
    logheader_t lh;
} log_t;
static log_t log;

static void recover_from_log();
static void commit();

static void log_reset() {
    log.outstanding = 0;
    log.committing = 0;
    log.seq = 0;
    log.ack = 0;
    recover_from_log();
}

void
initlog(int dev, struct superblock *sup)
{
    assert(sizeof(logheader_t) <= BSIZE);

    initlock(&log.lock, "log.lock");
    log.dev = dev;
    log.start = sup->logstart;
    log.size = sup->nlog;

    log_reset();

    cprintf("initlog: ok\n");
}

static void copy_log(int to_disk, int unpin) {
#ifdef LOG_TRACE
    cprintf("cpu %d: copy_log BEGIN: to_disk=%d, unpin=%d\n", arm_cpuid(), to_disk, unpin);
#endif

    for (int i = 0; i < log.lh.n; i++) {
        int si = log.lh.block[i];
        int di = log.start + 1 + i;
        assert(si != di);

        if (to_disk) {
            int tmp = si;
            si = di;
            di = tmp;
        }

        struct buf *src, *dst;
        src = bread(log.dev, si);
        dst = bread(log.dev, di);

        memcpy(dst->data, src->data, BSIZE);
        dst->flags |= B_DIRTY;
        bwrite(dst);

        if (unpin)
            b_unpin(to_disk ? dst : src);

        brelse(dst);
        brelse(src);
    }

#ifdef LOG_TRACE
    cprintf("cpu %d: copy_log END: to_disk=%d, unpin=%d\n", arm_cpuid(), to_disk, unpin);
#endif
}

/* Read the log header from disk into the in-memory log header. */
static void
read_head()
{
    struct buf *b = bread(log.dev, log.start);
    logheader_t *h = (logheader_t *) b->data;

    log.lh.n = h->n;
    for (int i = 0; i < LOGSIZE; i++) {
        log.lh.block[i] = h->block[i];
    }

    brelse(b);
}

/*
 * Write in-memory log header to disk.
 * This is the true point at which the
 * current transaction commits.
 */
static int write_head() {
#ifdef LOG_TRACE
    cprintf("cpu %d: write_head BEGIN\n", arm_cpuid());
#endif

    struct buf *buf = bread(log.dev, log.start);
    logheader_t *hb = (logheader_t *) (buf->data);
    int i;
    hb->n = log.lh.n;
    for (i = 0; i < log.lh.n; i++) {
        hb->block[i] = log.lh.block[i];
    }

    buf->flags |= B_DIRTY;
    int success = bwrite(buf);
    brelse(buf);

#ifdef LOG_TRACE
    cprintf("cpu %d: write_head END\n", arm_cpuid());
#endif

    return success;
}

static void
recover_from_log()
{
    read_head();

#ifdef LOG_TRACE
    cprintf("recover_from_log: found %d record(s).\n", log.lh.n);
#endif

    copy_log(1, 0);  // reinstall but do not unpin
    log.lh.n = 0;
    write_head();
}

static int min(int a, int b) {
    return a < b ? a : b;
}

/* Called at the start of each FS system call. */
void
begin_op()
{
    acquire(&log.lock);

    do {
        if (log.committing)
            sleep(&log, &log.lock);
        else if (log.lh.n + (log.outstanding + 1) * MAXOPBLOCKS > min(log.size - 1, LOGSIZE))
            sleep(&log, &log.lock);
        else {
            log.outstanding++;

#ifdef LOG_TRACE
            cprintf("cpu %d: begin_op: outstanding=%d\n", arm_cpuid(), log.outstanding);
#endif

            release(&log.lock);
            return;
        }
    } while (1);
}

/*
 * Called at the end of each FS system call.
 * Commits if this was the last outstanding operation.
 */
void
end_op()
{
    int do_commit = 0;

    acquire(&log.lock);

    assert(!log.committing);

#ifdef LOG_TRACE
    cprintf("cpu %d: end_op: outstanding=%d\n", arm_cpuid(), log.outstanding);
#endif

    log.outstanding--;
    log.seq++;
    if (log.outstanding == 0) {
        do_commit = 1;
        log.committing = 1;
    } else
        wakeup(&log);

    release(&log.lock);

    if (do_commit) {  // do not use log.committing since it's not protected by log.lock
        // since commit may acquire sleep-locks,
        // we should call it without holding log.lock.
        commit();

        acquire(&log.lock);
        log.committing = 0;
        wakeup(&log);
        release(&log.lock);
    }
}

static void
commit()
{
    if (log.lh.n > 0) {
        for (int i = 0; i < log.lh.n; i++) {
            struct buf *b = bread(log.dev, log.lh.block[i]);
            assert(b->refcnt == 2);
            brelse(b);
        }

        // write to log
        copy_log(0, 0);

        // write log header and then ack
        if (write_head()) {
            acquire(&log.lock);
            log.ack = log.seq;
            release(&log.lock);
        }

        // writeback to disk
        copy_log(1, 1);

        // clear log header
        acquire(&log.lock);
        log.lh.n = 0;
        release(&log.lock);

        write_head();
    }
}

/* Caller has modified b->data and is done with the buffer.
 * Record the block number and pin in the cache with B_DIRTY.
 * commit()/write_log() will do the disk write.
 *
 * log_write() replaces bwrite(); a typical use is:
 *   bp = bread(...)
 *   modify bp->data[]
 *   log_write(bp)
 *   brelse(bp)
 */
void
log_write(struct buf *b)
{
#ifdef LOG_TRACE
    cprintf("cpu %d: log_write: bno=%d\n", arm_cpuid(), b->blockno);
#endif

    acquire(&log.lock);

    assert(log.lh.n < LOGSIZE);
    assert(log.lh.n + 1 < log.size);  // +1 for header block
    assert(log.outstanding > 0);

    int i;
    for (i = 0; i < log.lh.n; i++) {
        if (log.lh.block[i] == b->blockno)
            break;
    }

    log.lh.block[i] = b->blockno;

    if (i == log.lh.n) {
        log.lh.n++;
        b_pin(b);
    }

    release(&log.lock);
}

// #define NROUND 100000
#define NROUND 100
#define NWORKER 8  // number of workers
#define NBLK 10     // number of blocks per transaction
#define BLK_ID(i) ((i) & 0xf)
#define NTX  16     // number of transactions per worker
#define TX_ID(i) (((i) >> 4) & 0xf)
#define WORKER_ID(i) ((i) >> 8)

#define SLIM 5000   // maximum number of syncs

static struct {
    struct spinlock lock;
    int remain, current;
    int data_start;
    int expected;
    char sb[BSIZE], lb[LOGSIZE + 1][BSIZE];
    int mem[NWORKER][NTX][NBLK];
    int valid[NWORKER][NTX];
} ctx;
static virtio_i backup;

#define BLKNO(wid, tid, bid) \
    (ctx.data_start + ((((wid) << 8) | (((tid) & 0xf)) << 4) | ((bid) & 0xf)))

static struct buf _sb_buf;

static void ctx_reset() {
    initlock(&ctx.lock, "log_test.ctx.lock");
    ctx.remain = SLIM;
    ctx.current = 0;
    ctx.data_start = log.start + log.size + 5;
    memset(ctx.lb, 0, sizeof(ctx.lb));
    memset(ctx.mem, 0, sizeof(ctx.mem));
    memset(ctx.valid, 0, sizeof(ctx.valid));

    _sb_buf.dev = 1;
    _sb_buf.blockno = 1;
    backup.sync(&_sb_buf);
    memcpy(ctx.sb, _sb_buf.data, BSIZE);
}

static int get_int_from_buf(struct buf *b) {
    return *((int *)b->data);
}

static void put_int_to_buf(struct buf *b, int v) {
    int *p = (int *)b->data;
    *p = v;
}

static int fake_sync(struct buf *b) {
    int bno = b->blockno;
    int valid = !!(b->flags & B_VALID);
    int dirty = !!(b->flags & B_DIRTY);

    int failed = 0;
    acquire(&ctx.lock);
    if (ctx.remain <= 0)
        failed = 1;
    else
        ctx.remain--;
    release(&ctx.lock);

    if (failed) {
#ifdef LOG_TRACE
        cprintf("cpu %d: fake_sync: bno=%d, valid=%d, dirty=%d (FAILED)\n",
            arm_cpuid(), b->blockno, valid, dirty);
#endif

        if (dirty) {
            /* do nothing */
        } else if (!valid)
            memset(b->data, 0xcc, BSIZE);
    } else if (bno == 1 || (log.start <= bno && bno < log.start + log.size)) {
        char *p = bno == 1 ? ctx.sb : (ctx.lb[bno - log.start] + 0);

#ifdef LOG_TRACE
        cprintf("cpu %d: fake_sync: bno=%d, valid=%d, dirty=%d\n",
            arm_cpuid(), b->blockno, valid, dirty);
#endif

        if (dirty)
            memcpy(p, b->data, BSIZE);
        else if (!valid)
            memcpy(b->data, p, BSIZE);
    } else if (bno >= ctx.data_start) {
        bno -= ctx.data_start;

        int wid, tid, bid;
        wid = WORKER_ID(bno);
        tid = TX_ID(bno);
        bid = BLK_ID(bno);

#ifdef LOG_TRACE
        cprintf("cpu %d: fake_sync: bno=%d, valid=%d, dirty=%d (wid=%d, tid=%d, bid=%d)\n",
            arm_cpuid(), b->blockno, valid, dirty, wid, tid, bid);
#endif

        if (dirty)
            ctx.mem[wid][tid][bid] = get_int_from_buf(b);
        else if (!valid)
            put_int_to_buf(b, ctx.mem[wid][tid][bid]);
    } else
        panic("unknown bno=%d", bno);

    b->flags = B_VALID;
    return !failed;
}

static void log_test_worker() {
    acquire(&ctx.lock);
    int wid = ctx.current++;
    release(&ctx.lock);

    for (int tid = 0; tid < NTX; tid++) {
        begin_op();

        for (int bid = 0; bid < NBLK; /* bid++ */) {
            int v;
            if (tid == 0)
                v = 1;
            else {
                int i;

                acquire(&ctx.lock);
                for (i = 0; i < NWORKER; i++) {
                    if (ctx.valid[i][tid - 1])
                        break;
                }
                release(&ctx.lock);

                assert(i < NWORKER);

                struct buf *b = bread(log.dev, BLKNO(wid, tid - 1, bid));
                v = get_int_from_buf(b) + 1;
                brelse(b);
            }

            struct buf *b = bread(log.dev, BLKNO(wid, tid, bid));
            put_int_to_buf(b, v);
            b->flags |= B_DIRTY;
            log_write(b);
            assert(b->refcnt == 2);
            brelse(b);

            // increment with a probability of 3/4
            if (xrand() & 0x3)
                bid++;
        }

        end_op();

        acquire(&ctx.lock);
        ctx.valid[wid][tid] = 1;
        if (ctx.remain <= 0) {  // early return after end_op
            release(&ctx.lock);
            return;
        }
        release(&ctx.lock);
    }
}

static void log_test_check() {
    assert(log.lh.n == 0);

#ifdef LOG_TRACE
    cprintf("log_test_check: expect %d tx.\n", ctx.expected);
#endif

    int tx_count = 0;
    for (int wid = 0; wid < NWORKER; wid++) {
        for (int tid = 0; tid < NTX; tid++) {
            int *w = ctx.mem[wid][tid] + 0;
            int asum = -1, osum = 0;
            for (int bid = 0; bid < NBLK; bid++) {
                asum &= w[bid];
                osum |= w[bid];
            }

#ifdef LOG_TRACE
            cprintf("log_test_check: wid=%d, tid=%d: and=%d, or=%d\n",
                wid, tid, asum, osum);
#endif

            assert(asum == osum);
            assert(asum == tid + 1 || asum == 0);

            if (asum != 0)
                tx_count++;
        }
    }

    assert(tx_count == ctx.expected);
}

void log_test() {
    cprintf("log_test: logging test started.\n");

    // hook virtio
    b_invalidate();
    backup = virtio;
    virtio.sync = fake_sync;

    // run tests
    for (int t = 0; t < NROUND; t++) {
        cprintf("log_test: round %d...", t);

        b_invalidate();

        ctx_reset();
        task_run(NWORKER, log_test_worker);

#ifdef LOG_TRACE
        cprintf("log_test: reset\n");
#endif

        ctx.expected = log.ack;
        ctx.remain = INT32_MAX;
        b_full_reset();
        initlog(log.dev, (struct superblock *)ctx.sb);
        log_test_check();
    }

    // recover virtio
    b_invalidate();
    virtio = backup;

    cprintf("log_test: logging test passed.\n");
}